import os
import chromadb
from llama_index.core import Settings, SimpleDirectoryReader, VectorStoreIndex
from llama_index.core.node_parser import SimpleNodeParser
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.huggingface import HuggingFaceLLM
from llama_index.vector_stores.chroma import ChromaVectorStore
from modelscope import snapshot_download

# ===========================
# 一、环境准备
# ===========================
# 设置模型缓存路径（避免重复下载）
os.environ['MODELSCOPE_CACHE'] = r"D:\models"

# ===========================
# 二、加载本地大语言模型（LLM）
# ===========================
# 从 ModelScope 下载 deepseek 的本地模型
llm_model_dir = snapshot_download("deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B")

# 初始化 LLM（本地推理用）
llm = HuggingFaceLLM(
    model_name=llm_model_dir,       # 模型路径
    tokenizer_name=llm_model_dir,   # 分词器路径（通常与模型相同）
    context_window=2048,            # 最大上下文长度（输入 token 限制）
    max_new_tokens=256,             # 模型回答时生成的最大 token 数
    model_kwargs={
        "trust_remote_code": True,  # 某些 HuggingFace 模型需要这个参数
        "torch_dtype": "auto"       # 自动选择计算精度（fp16/fp32）
    },
    tokenizer_kwargs={
        "trust_remote_code": True
    },
    generate_kwargs={
        "temperature": 0.7,         # 控制输出的随机性，越高越随机
        "do_sample": True           # 启用采样
    }
)

# ===========================
# 三、加载向量化模型（Embedding）
# ===========================
embedding_model_dir = snapshot_download("sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2")

# 初始化文本向量化模型，用于将文本转为向量
embedding_model = HuggingFaceEmbedding(
    model_name=embedding_model_dir
)

# ===========================
# 四、注册全局配置
# ===========================
Settings.llm = llm
Settings.embed_model = embedding_model


# ===========================
# 六、加载本地文档数据
# ===========================
# 从 data 文件夹读取文本数据（支持 .txt、.md、.pdf 等）
documents = SimpleDirectoryReader(r"data").load_data()

# ===========================
# 七、分块处理文档（切片成更小的片段）
# ===========================
# 每个 chunk 大约 512 个 token，便于检索与匹配
node_parser = SimpleNodeParser.from_defaults(chunk_size=512)

# 将文档转为节点（Node 是 llamaindex 内部的数据单元）
base_nodes = node_parser.get_nodes_from_documents(documents)

print("加载到的文档：", documents)
print("==================================")
print("生成的节点：", base_nodes)

# ===========================
# 八、创建索引（将节点写入向量数据库）
# ===========================
index = VectorStoreIndex(nodes=base_nodes)

# 将索引数据保存到磁盘（持久化存储）
index.storage_context.persist()

# ===========================
# 九、查询阶段（RAG 查询）
# ===========================
# 将索引转换为查询引擎
query_engine = index.as_query_engine()

# 输入问题，模型会基于检索的文档内容生成回答
response = query_engine.query("王德宝有多少钱")

print("模型回答：")
print(response)
